from typing import Any

from datasets import Dataset as HFDataset, IterableDataset as HFIterable

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator


class HuggingFaceDatasetIterator(DatasetIterator[Any]):
    def __init__(self, dataset: HFDataset | HFIterable, state: int | None = None):
        self._dataset = dataset
        self._index = state if state is not None else 0
        self._iterator = iter(self._dataset.skip(self._index))

    def __next__(self) -> Any:
        sample = next(self._iterator)
        self._index += 1
        return sample

    def __skip__(self, samples_n: int):
        self._index += samples_n
        self._iterator = iter(self._dataset.skip(self._index))

    def __pause__(self) -> int:
        return self._index


class HuggingFaceDataset(Dataset[Any]):
    def __init__(self, dataset: HFDataset | HFIterable):
        self._dataset = dataset

    @property
    def dataset(self) -> HFDataset | HFIterable:
        return self._dataset

    def __len__(self) -> int:
        return len(self._dataset)

    def __getitem__(self, index: int) -> Any:
        return self._dataset[index]

    def __iter__(self) -> HuggingFaceDatasetIterator:
        return HuggingFaceDatasetIterator(self._dataset)

    def __resume__(self, state: int | None = None, /) -> HuggingFaceDatasetIterator:
        return HuggingFaceDatasetIterator(self._dataset, state)
